#include "cybertron/network/SocketUdpReceiver.hpp"
#include "cybertron/network/SocketUdpCommon.hpp"
#include "cybertron/core/UtilCRT.hpp"
#include "cybertron/core/Log.hpp"
#include <thread>
#include <iostream>
#include <unordered_map>
#include <asio.hpp>
#include "cybertron/network/perf.hpp"

CYBERTRON_BEGIN

class SocketUdpReceiverImpl
{
public:
	std::mutex															mCachedMessageMutex;
	std::list<std::string>												mCachedMessages;
	std::list<std::shared_ptr<Message> >								mCachedCybertronMessages;
	std::mutex															mUDPMessageFragmentMutex;
	std::list<UDPMessageFragment>										mUDPMessageFragments;
	std::mutex															mUDPMessageProcessingMutex;
	std::unordered_map<std::uint16_t, std::shared_ptr<UDPMessage> >		mUDPMessagesProcessing;
	bool mbSupportMessageSplit;

	std::shared_ptr<asio::ip::udp::socket> mpSocket;
	std::shared_ptr<asio::io_service> mpIoService;
	std::shared_ptr<bool> mThreadRunning;
	std::shared_ptr<std::thread> mpSenderThread;
	std::shared_ptr<std::thread> mpProcessingThread;
	int mNumMaximalMessageCache = -1;
};

SocketUdpReceiver::SocketUdpReceiver(bool bSupportMessageSplit/* = true*/)
{
	mpImpl = new SocketUdpReceiverImpl();
	mpImpl->mbSupportMessageSplit = bSupportMessageSplit;
}

SocketUdpReceiver::~SocketUdpReceiver()
{
	clear();
	delete mpImpl;
	mpImpl = nullptr;
}

bool SocketUdpReceiver::initialize(
	int port,
	int numMaximalMessageCache)
{
	clear();

	mpImpl->mNumMaximalMessageCache = numMaximalMessageCache;
	mpImpl->mThreadRunning = std::make_shared<bool>(true);
	mpImpl->mpIoService = std::make_shared<asio::io_service>();
	mpImpl->mpSocket = std::make_shared<asio::ip::udp::socket>(
		*(mpImpl->mpIoService),
		asio::ip::udp::endpoint(asio::ip::udp::v4(), port));
#if CYBERTRON_WIN || CYBERTRON_LINUX
	asio::socket_base::receive_buffer_size option(UDP_MAX_MESSAGE_SIZE);
	mpImpl->mpSocket->set_option(option);
	mpImpl->mpSocket->get_option(option);
	logInfo("Receive Buffer Size: %d\r\n", option.value());
#endif

	// start background thread
	auto pThreadRunning = mpImpl->mThreadRunning;

	mpImpl->mpSenderThread = UtilCRT::startThread(
		"UdpReceiver",
		[pThreadRunning, numMaximalMessageCache, this]()
	{
		std::array<char, UDP_MAX_FRAGMENT_SIZE> buff;
		while (*pThreadRunning)
		{
			asio::ip::udp::endpoint senderEndPoint;
			std::size_t bytesReceived =
				mpImpl->mpSocket->receive_from(asio::buffer(buff), senderEndPoint);

			IoPerf::inst().incr_recv_packs(1);
			IoPerf::inst().incr_recv_bytes(bytesReceived);

			if (bytesReceived < UDP_MIN_MESSAGE_SIZE) {
				continue;
			}
			if (mpImpl->mbSupportMessageSplit) {
				UDPMessageFragment messageFragment;
				if (!SocketUdpCommon::decodeMessage(buff, (std::uint16_t)bytesReceived, messageFragment)) {
					logError("Invalid message received!");
					continue;
				}
				std::lock_guard<std::mutex> lock(mpImpl->mUDPMessageFragmentMutex);
				mpImpl->mUDPMessageFragments.emplace_back(std::move(messageFragment));
			}
			else {
				std::string message = SocketUdpCommon::decodeSimpleMessage(buff, (std::uint16_t)bytesReceived);
				if (message.empty()) {
					logError("Invalid message received!");
					continue;
				}
				//logDebug("Received Message Size : %d", message.size());

				std::lock_guard<std::mutex> lock(mpImpl->mCachedMessageMutex);
				mpImpl->mCachedMessages.emplace_back(std::move(message));
				if (numMaximalMessageCache > 0 && mpImpl->mCachedMessages.size() > (size_t)numMaximalMessageCache) {
					mpImpl->mCachedMessages.pop_front();
				}
			}
		}
	});

	mpImpl->mpProcessingThread = UtilCRT::startThread(
		"UdpMessageProcessor",
		[pThreadRunning, numMaximalMessageCache, this]()
	{
		while (*pThreadRunning)
		{
			std::list<UDPMessageFragment> udpMessageFragments;
			{
				std::lock_guard<std::mutex> lock(mpImpl->mUDPMessageFragmentMutex);
				udpMessageFragments.swap(mpImpl->mUDPMessageFragments);
			}
			if (udpMessageFragments.empty()) {
				// no new message wait for 2ms
				std::this_thread::sleep_for(std::chrono::milliseconds(2));
				continue;
			}
			std::list<std::shared_ptr<UDPMessage> > readyMessages;
			{
				std::lock_guard<std::mutex> lock(mpImpl->mUDPMessageProcessingMutex);
				// clear overdue first.
				std::list<std::uint16_t> overdueMessageIds;
				for (const auto& keyValue : mpImpl->mUDPMessagesProcessing) {
					if (keyValue.second->isMessageOverdue()) {
						overdueMessageIds.push_back(keyValue.first);
					}
				}
				for (auto overdueMessageId : overdueMessageIds) {
					mpImpl->mUDPMessagesProcessing.erase(overdueMessageId);
					logWarning("Message overdue for id : %d\r\n\r\n", overdueMessageId);
				}
				// process new
				for (auto& messageFragment : udpMessageFragments) {
					auto seqId = messageFragment.SequenceId;
					auto iter = mpImpl->mUDPMessagesProcessing.find(seqId);
					if (iter == mpImpl->mUDPMessagesProcessing.end()) {
						auto pMessage = std::make_shared<UDPMessage>();
						pMessage->addMessageFragment(messageFragment);
						if (pMessage->isMessageReady()) {
							readyMessages.emplace_back(pMessage);
						}
						else {
							mpImpl->mUDPMessagesProcessing.emplace(seqId, pMessage);
						}
					}
					else {
						auto pExistingMessage = iter->second;
						pExistingMessage->addMessageFragment(messageFragment);
						if (pExistingMessage->isMessageReady()) {
							readyMessages.emplace_back(pExistingMessage);
							mpImpl->mUDPMessagesProcessing.erase(iter);
						}
					}
				}
				udpMessageFragments.clear();
			}
			
			{
				std::lock_guard<std::mutex> lock(mpImpl->mCachedMessageMutex);
				for (auto& udpMessage : readyMessages) {
					std::string strMessage = udpMessage->decodeMessage();
					if (strMessage.empty()) {
						logError("Invalid message in decode!");
						continue;
					}
					// logDebug("Decoded Message Size : %d", message.size());
					auto pMessage = std::make_shared<Message>();
					pMessage->create(0, 0);
					memcpy(pMessage->getData(), &strMessage[0], pMessage->getLength());
					auto dataBodyLength = pMessage->parseDataBodyLength();
					if (dataBodyLength + pMessage->getLength() != strMessage.size()) {
						// logError("Invalid message found when converting back to Message!");
						continue;
					}
					pMessage->create(pMessage->parseMsgId(), dataBodyLength);
					memcpy(pMessage->getData(), &strMessage[0], pMessage->getLength());
					mpImpl->mCachedCybertronMessages.emplace_back(pMessage);
					if (numMaximalMessageCache > 0 && mpImpl->mCachedCybertronMessages.size() > (size_t)numMaximalMessageCache) {
						mpImpl->mCachedCybertronMessages.pop_front();
					}
				}
				readyMessages.clear();
			}
		}
	});

	return true;
}

void SocketUdpReceiver::clear()
{
	if (mpImpl->mThreadRunning != nullptr)
	{
		(*mpImpl->mThreadRunning) = false;

		if (mpImpl->mpIoService != nullptr)
		{
			mpImpl->mpIoService->stop();
		}
	}

	if (mpImpl->mpSocket != nullptr)
	{
		mpImpl->mpSocket->close();
		mpImpl->mpSocket = nullptr;
	}

	if (mpImpl->mpSenderThread != nullptr)
	{
		mpImpl->mpSenderThread->join();
		mpImpl->mpSenderThread = nullptr;
	}

	if (mpImpl->mpProcessingThread != nullptr) {
		mpImpl->mpProcessingThread->join();
		mpImpl->mpProcessingThread = nullptr;
	}
	mpImpl->mThreadRunning = nullptr;
	mpImpl->mpIoService = nullptr;
	mpImpl->mNumMaximalMessageCache = -1;

	{
		std::lock_guard<std::mutex> lock(mpImpl->mUDPMessageProcessingMutex);
		mpImpl->mUDPMessagesProcessing.clear();
	}
	{
		std::lock_guard<std::mutex> lock(mpImpl->mUDPMessageFragmentMutex);
		mpImpl->mUDPMessageFragments.clear();
	}
	{
		std::lock_guard<std::mutex> lock(mpImpl->mCachedMessageMutex);
		mpImpl->mCachedMessages.clear();
		mpImpl->mCachedCybertronMessages.clear();
	}
	
}

void SocketUdpReceiver::takeAwayMessages(
	std::list<std::string>& messages)
{
	messages.clear();
	std::lock_guard<std::mutex> lock(mpImpl->mCachedMessageMutex);
	mpImpl->mCachedMessages.swap(messages);
}

void SocketUdpReceiver::takeAwayMessages(std::list<std::shared_ptr<Message> >& messages) {
	messages.clear();
	std::lock_guard<std::mutex> lock(mpImpl->mCachedMessageMutex);
	mpImpl->mCachedCybertronMessages.swap(messages);
}

CYBERTRON_END
